# 导入必要的库
# numpy: 用于处理数组和数值计算
# sklearn.tree.DecisionTreeClassifier: 决策树分类器
# sklearn.svm.SVC: 支持向量机分类器
# sklearn.metrics: 用于计算模型性能指标
# time: 用于记录训练时间
# matplotlib.pyplot: 用于绘制图表
# matplotlib.rcParams: 用于设置图表字体和样式
# sklearn.base.clone: 用于复制分类器
# os: 用于文件操作
# datetime: 用于生成时间戳
import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, f1_score
import time
import matplotlib.pyplot as plt
from matplotlib import rcParams
from sklearn.base import clone
import os
from datetime import datetime

# 设置 Matplotlib 的字体为支持中文的字体
rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 加载预处理好的MNIST数据集
# 数据集包含训练集和测试集的特征和标签
print("加载预处理好的MNIST数据集...")
data = np.load('./mnist_processed.npz')
X_train = data['X_train']  # 训练集特征
X_test = data['X_test']    # 测试集特征
y_train = data['y_train']  # 训练集标签
y_test = data['y_test']    # 测试集标签

# 打印数据集信息
# 包括训练集大小、测试集大小和类别数量
print("\n数据集信息:")
print(f"训练集大小: {X_train.shape}")  # 训练集的样本数量和特征维度
print(f"测试集大小: {X_test.shape}")   # 测试集的样本数量和特征维度
print(f"类别数量: {len(np.unique(y_train))}") # 数据集中的类别数量

# 定义AdaBoost类，用于实现AdaBoost算法
# 支持多分类任务
class AdaBoost:
    # 初始化函数
    # base_estimator: 基分类器
    # n_estimators: 基分类器数量
    def __init__(self, base_estimator, n_estimators=50):
        self.base_estimator = base_estimator
        self.n_estimators = n_estimators
        self.alphas = []  # 存储每个分类器的权重
        self.estimators = []  # 存储基分类器
        self.classes = None  # 存储类别信息
        self.errors = []  # 存储每次迭代的错误率
    
    # 训练AdaBoost模型
    def fit(self, X, y):
        n_samples = X.shape[0]
        self.classes = np.unique(y)  # 获取类别信息
        n_classes = len(self.classes)
        
        # 初始化样本权重
        w = np.ones(n_samples) / n_samples
        
        for i in range(self.n_estimators):
            estimator = clone(self.base_estimator)  # 克隆基分类器
            
            # 特殊处理SVM（移除概率计算）
            if isinstance(estimator, SVC):
                estimator.probability = False
                
            estimator.fit(X, y, sample_weight=w)  # 使用样本权重训练分类器
            y_pred = estimator.predict(X)  # 预测训练集
            
            incorrect = (y_pred != y)  # 标记预测错误的样本
            err = np.sum(w * incorrect) + 1e-10  # 计算错误率
            
            # 早停机制：如果错误率过高，停止训练
            if err > 1-1/n_classes:
                print(f"迭代 {i+1}: 错误率{err:.4f}过大, 提前终止")
                break
                
            # SAMME算法计算分类器权重alpha，多加的一部分是SAMME
            alpha = np.log((1 - err) / err) + np.log(n_classes - 1)
            
            # 更新样本权重
            # 13-52 的方法是 w *= np.exp(-alpha * (y_pred == y)) 没有很大区别
            # w *= np.exp(alpha * incorrect - alpha * (y_pred == y)) 的方法尝试了，差很多，虽然它让alpha上升但实际效果欠佳
            w *= np.exp(alpha * incorrect)
            w /= np.sum(w)  # 归一化
            
            # 保存模型参数
            self.alphas.append(alpha)
            self.estimators.append(estimator)
            self.errors.append(err)
            
            # 打印进度
            print(f"迭代 {i+1}/{self.n_estimators}: 添加{type(estimator).__name__}分类器, 错误率={err:.4f}, alpha={alpha:.4f}")
    
    # 使用所有基分类器进行预测
    def predict(self, X):
        return self._partial_predict(X, len(self.estimators))
    
    # 使用前n个基分类器进行预测
    def _partial_predict(self, X, n_classifiers):
        n_classes = len(self.classes)
        class_scores = np.zeros((X.shape[0], n_classes))  # 初始化类别得分
        
        for i in range(n_classifiers):
            alpha = self.alphas[i]
            estimator = self.estimators[i]
            y_pred = estimator.predict(X)
            
            # 为每个样本的预测类别增加alpha权重
            for j, cls in enumerate(self.classes):
                class_scores[:, j] += alpha * (y_pred == cls)
        
        # 返回得分最高的类别
        return self.classes[np.argmax(class_scores, axis=1)]
    
    # 可视化基分类器数量对性能的影响
    def plot_classifier_impact(self, X, y, title="基分类器数量影响（优化版）", save_dir=None):
        """可视化基分类器数量对性能的影响（优化版）"""
        n = len(self.estimators)
        
        # 预计算所有基分类器的预测结果
        all_preds = []
        print(f"预计算所有基分类器的预测...")
        for i, estimator in enumerate(self.estimators):
            # 显示进度
            if (i+1) % 5 == 0 or i == 0 or i == n-1:
                print(f"预计算基分类器 {i+1}/{n}...")
            
            # 存储预测结果
            all_preds.append(estimator.predict(X))
        
        # 设置采样点
        step = max(1, n // 20)  # 均匀分布20个点
        points = list(range(1, n+1, step))
        if points[-1] != n:
            points.append(n)  # 确保包含最后一个点
        
        accuracies = []
        
        # 计算每个点的准确率
        print(f"计算不同基分类器数量的准确率...")
        for i, num_clf in enumerate(points):
            # 显示进度
            print(f"处理 {i+1}/{len(points)}: {num_clf}个基分类器")
            
            # 初始化类别得分
            n_classes = len(self.classes)
            class_scores = np.zeros((X.shape[0], n_classes))
            
            # 累加前 num_clf 个分类器的预测
            for j in range(num_clf):
                y_pred = all_preds[j]  # 使用预计算的预测
                alpha = self.alphas[j]
                
                # 更新类别得分
                for k, cls in enumerate(self.classes):
                    class_scores[:, k] += alpha * (y_pred == cls)
            
            # 预测最终标签
            y_pred_final = self.classes[np.argmax(class_scores, axis=1)]
            acc = accuracy_score(y, y_pred_final)
            accuracies.append(acc)
        
        # 绘图
        plt.figure(figsize=(10, 6))
        plt.plot(points, accuracies, 'b-o')
        plt.title(title)
        plt.xlabel('基分类器数量')
        plt.ylabel('测试准确率')
        plt.grid(True)
        
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)
            plot_path = os.path.join(save_dir, f"{title.replace(' ', '_')}.png")
            plt.savefig(plot_path)  # 确保保存图像
            print(f"图表已保存为 '{plot_path}'")
        
        return accuracies

# 从训练集中随机抽取样本，默认取10000个
# 用于减少训练时间
def sample_training_data(X_train, y_train, sample_size=10000):
    indices = np.random.choice(X_train.shape[0], sample_size, replace=False)
    return X_train[indices], y_train[indices]

# 运行AdaBoost性能对比实验
# 比较不同基分类器的性能
def run_adaboost_experiment():
    print("\n" + "="*50)
    print("AdaBoost性能对比实验")
    print("="*50)
    
    results = {}
    
    # 获取当前时间戳
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    
    # 创建结果保存目录
    save_dir = './ada_results'
    os.makedirs(save_dir, exist_ok=True)
    
    # 决策树桩基分类器 (最大深度=1)
    print("\n训练AdaBoost (决策树桩基分类器)...")
    start_time = time.time()
    
    # 抽取样本
    X_train_sampled, y_train_sampled = sample_training_data(X_train, y_train, sample_size=10000)
    
    adaboost_dt = AdaBoost(
        base_estimator=DecisionTreeClassifier(max_depth=1),  # 决策树桩，最大深度为1
        n_estimators=500  # 基分类器数量，设置为500以观察性能变化
    )
    adaboost_dt.fit(X_train_sampled, y_train_sampled)
    dt_train_time = time.time() - start_time
    
    # 评估性能
    y_pred_dt = adaboost_dt.predict(X_test)
    dt_acc = accuracy_score(y_test, y_pred_dt)
    dt_f1 = f1_score(y_test, y_pred_dt, average='weighted')
    
    # 保存预测结果
    pred_path_dt = os.path.join(save_dir, f'adaboost_dt_predictions_{timestamp}.txt')
    np.savetxt(pred_path_dt, y_pred_dt[:10], fmt='%d')  # 仅保存前十个
    print(f"决策树桩预测结果已保存为 '{pred_path_dt}'")
    
    # 绘制分类器数量影响图
    print("绘制决策树桩AdaBoost分类器数量影响图...")
    adaboost_dt.plot_classifier_impact(X_test, y_test, f'AdaBoost (决策树桩) 分类器数量影响 {timestamp}', save_dir)
    
    # 线性SVM基分类器 (减少迭代次数)
    print("\n训练AdaBoost (线性SVM基分类器)...")
    start_time = time.time()
    
    # 抽取样本
    X_train_sampled, y_train_sampled = sample_training_data(X_train, y_train, sample_size=10000)
    
    adaboost_svm = AdaBoost(
        base_estimator=SVC(
            kernel='linear',  # 使用线性核函数
            random_state=42,  # 随机种子，确保结果可复现
            decision_function_shape='ovr'  # 一对多策略，用于多分类任务
        ),
        n_estimators=20  # 基分类器数量，设置为20以减少训练时间
    )
    adaboost_svm.fit(X_train_sampled, y_train_sampled)
    svm_train_time = time.time() - start_time
    
    # 评估性能
    y_pred_svm = adaboost_svm.predict(X_test)
    svm_acc = accuracy_score(y_test, y_pred_svm)
    svm_f1 = f1_score(y_test, y_pred_svm, average='weighted')
    
    # 保存预测结果
    pred_path_svm = os.path.join(save_dir, f'adaboost_svm_predictions_{timestamp}.txt')
    np.savetxt(pred_path_svm, y_pred_svm[:10], fmt='%d')  # 仅保存前十个
    print(f"SVM预测结果已保存为 '{pred_path_svm}'")
    
    # 绘制分类器数量影响图
    print("绘制SVM AdaBoost分类器数量影响图...")
    adaboost_svm.plot_classifier_impact(X_test, y_test, f'AdaBoost (线性SVM) 分类器数量影响 {timestamp}', save_dir)
    
    # 保存综合结果日志
    log_path = os.path.join(save_dir, f'adaboost_results_{timestamp}.txt')
    with open(log_path, 'w') as f:
        f.write(f"AdaBoost性能对比 ({timestamp}):\n")
        f.write("决策树桩基分类器:\n")
        f.write(f"  准确率: {dt_acc:.4f}, F1分数: {dt_f1:.4f}, 训练时间: {dt_train_time:.2f}秒\n")
        f.write("\n线性SVM基分类器:\n")
        f.write(f"  准确率: {svm_acc:.4f}, F1分数: {svm_f1:.4f}, 训练时间: {svm_train_time:.2f}秒\n")
    print(f"综合结果日志已保存为 '{log_path}'")
    
    results['DecisionTree'] = {
        'accuracy': dt_acc,
        'f1_score': dt_f1,
        'train_time': dt_train_time
    }
    
    results['LinearSVM'] = {
        'accuracy': svm_acc,
        'f1_score': svm_f1,
        'train_time': svm_train_time
    }
    
    return results

# 主函数
# 运行AdaBoost实验并打印总结
def main():
    # 运行AdaBoost实验
    adaboost_results = run_adaboost_experiment()
    
    # 打印最终总结
    print("\n" + "="*50)
    print("实验总结")
    print("="*50)
    
    print("\nAdaBoost性能对比:")
    for base, metrics in adaboost_results.items():
        print(f"{base}: 准确率={metrics['accuracy']:.4f}, F1分数={metrics['f1_score']:.4f}, 训练时间={metrics['train_time']:.2f}秒")

if __name__ == "__main__":
    main()